/*
 * EquilibriumSolver.cpp
 *
 *  Created on: 26 Jul 2011
 *      Author: Allan
 */

#include "EquilibriumSolver.h"

// C++ includes
#include <cmath>
#include <functional>
#include <iomanip>
#include <memory>
using namespace std;
using namespace std::placeholders;

// Boost includes
#include <boost/foreach.hpp>
using namespace boost;

// GeoReact includes
#include "EquilibriumProblem.h"
#include "EquilibriumProblemReduced.h"
#include "../MultirootSolver.h"
#include "../SpeciesMap.h"
#include "../Utils.h"

EquilibriumSolver::EquilibriumSolver() :
tolerance(1.0E-6), maxiterations(100), reducedFormulation(false)
{}

EquilibriumSolver::EquilibriumSolver(const Multiphase& multiphase, const ReactionSystem& equilibriumReactions) :
tolerance(1.0E-6), maxiterations(100), reducedFormulation(false)
{
	Initialiaze(multiphase, equilibriumReactions);
}

void EquilibriumSolver::Initialiaze(const Multiphase& multiphase, const ReactionSystem& equilibriumReactions)
{
	// Set the phase system object
	this->multiphase = multiphase;
	
	// Set the equilibrium reaction system object
	this->equilibriumReactions = equilibriumReactions;
			
	// The species map object to determine the indexes of the species
	const SpeciesMap& speciesMap = multiphase.GetSpeciesMap();
	
	// Set the indexes data of the system of equilibrium-controlled reactions
	this->equilibriumReactions.SetIndexes(speciesMap);
	
	// Set the indexes of the primary and equilibrium species
	jIndexes = speciesMap[ equilibriumReactions.Products() ];
	eIndexes = speciesMap[ equilibriumReactions.Reactants() ];
	
	// Set the number of primary (Nj) and equilibrium (Ne) species
	Nj = jIndexes.size();
	Ne = eIndexes.size();
	
	// Set the stoichiometric matrix of the equilibrium-controlled reactions
	vej = equilibriumReactions.StoichiometricMatrix();
	
	// Initialise the Eigen matrices with zero entries
	alpha = MatrixXd::Zero(Nj, Nj);
	 beta = MatrixXd::Zero(Nj, Ne);
	   
	// Resize the vector of condition functions
	hfunctions.resize(Nj);
}

void EquilibriumSolver::SetEquilibriumConditions(const vector<EquilibriumCondition>& conditions)
{
	// Loop over all the equilibrium conditions
	for(uint i = 0; i < conditions.size(); ++i)
	{
		// Set the i-th row of the alpha matrix
		alpha.row(i) = conditions[i].alphaRow;
		
		// Set the i-th row of the beta matrix
		beta.row(i) = conditions[i].betaRow;
		
		// Set the i-th row of the equilibrium condition functions
		hfunctions[i] = conditions[i].h;
	}
}

void EquilibriumSolver::SetTotalCompositionMode()
{
	// Set the (alpha) matrix
	alpha.setIdentity(Nj, Nj);
	
	// Set the (beta) matrix
	beta = vej.transpose();
}

void EquilibriumSolver::SetTotalComposition(const VectorXd& uj)
{
	for(uint j = 0; j < Nj; ++j) hfunctions[j] = Constant(uj[j]);
}

void EquilibriumSolver::SetSolverOptions(double tolerance, uint maxiterations)
{
	this->tolerance = tolerance;
	this->maxiterations = maxiterations;
}

void EquilibriumSolver::UseReducedFormulation()
{
	reducedFormulation = true;
}

const EquilibriumSolver::Diagnostic EquilibriumSolver::Solve(double T, double P, VectorXd& n) const
{
	// Check which formulation to be used
	if(reducedFormulation)
		return SolveReducedFormulation(T, P, n);
	else
		return SolveGeneralFormulation(T, P, n);
}

const EquilibriumSolver::Diagnostic EquilibriumSolver::SolveGeneralFormulation(double T, double P, VectorXd& n) const
{
	// The vector of moles of primary/equilibrium species (njne) = [nj ne]
	VectorXd njne(Nj + Ne);
	
	// Initialise the vector (njne) with the values from the initial guess (nGuess)
	for(uint j = 0; j < Nj; ++j) njne[j]      = n[jIndexes[j]];
	for(uint e = 0; e < Ne; ++e) njne[e + Nj] = n[eIndexes[e]];
	
	// Create an instance of the definition of the equilibrium problem
	EquilibriumProblem problem(*this, T, P, n);
	
	// Create an instance of the multiroot solver
	MultirootSolver solver(Nj + Ne);
	
	// Set the algorithm of the solver to the modified Newton-Raphson method (the standard Newton-Raphson algorithm works too, but the algorithms Hybridj and Hybridsj does not) 
	solver.SetAlgorithm(MultirootSolver::StandardNewton);
	
	// Set the Newton solver options
	solver.SetSolverOptions(tolerance, maxiterations);
	
	// Set the diagnostic instance
	Diagnostic diagnostic = {0, false};
	
	// The solution vector (njne) at the previous iteration (used for application of under relaxation to the solution)
	VectorXd njneOld = njne;
	
	// Solve the equilibrium problem
	while(solver.Iterate(problem, njne) && diagnostic.iterations < maxiterations)
	{
		// Apply the under relaxation to the vector (njne) to avoid negative values of compositions
		UnderRelax(njne, njneOld);
		
		// Update the vector (njneOld)
		njneOld = njne;
		
		// Update (iteration)
		++diagnostic.iterations;
	}
	
	// Check convergence of the previous iterations
	diagnostic.converged = diagnostic.iterations < maxiterations;
	
	// Transfer the solution (njne) to the vector (n)
	for(uint j = 0; j < Nj; ++j) n[jIndexes[j]] = njne[j];
	for(uint e = 0; e < Ne; ++e) n[eIndexes[e]] = njne[e + Nj];
	
	return diagnostic;
}
	
const EquilibriumSolver::Diagnostic EquilibriumSolver::SolveReducedFormulation(double T, double P, VectorXd& n) const
{
	// The vector of moles of primary species (nj)
	VectorXd nj(Nj);
	
	// Initialise the vector (nj) with the values from the initial guess (n)
	for(uint j = 0; j < Nj; ++j) nj[j] = n[jIndexes[j]];
	
	// Create an instance of the definition of the equilibrium problem
	EquilibriumProblemReduced problem(*this, T, P, n);
	
	// Create an instance of the multiroot solver
	MultirootSolver solver(Nj);
	
	// Set the algorithm of the solver to the modified Newton-Raphson method (the standard Newton-Raphson algorithm works too, but the algorithms Hybridj and Hybridsj does not) 
	solver.SetAlgorithm(MultirootSolver::StandardNewton);
	
	// Set the Newton solver options
	solver.SetSolverOptions(tolerance, maxiterations);
	
	// Set the diagnostic instance
	Diagnostic diagnostic = {0, false};
	
	// The solution vector (nj) at the previous iteration (used for application of under relaxation to the solution)
	VectorXd njOld = nj;
	
	// Solve the equilibrium problem
	while(solver.Iterate(problem, nj) && diagnostic.iterations < maxiterations)
	{
		// Apply the under relaxation to the vector (nj) to avoid negative values of compositions
		UnderRelax(nj, njOld);
		
		// Update the vector (njOld)
		njOld = nj;
		
		// Update (iteration)
		++diagnostic.iterations;
	}
	
	// Check convergence of the previous iterations
	diagnostic.converged = diagnostic.iterations < maxiterations;
	
	// Transfer the solution (nj) and (ne) to the vector (n) (use the equilibrium problem instance to get the (ne) vector)
	for(uint j = 0; j < Nj; ++j) n[jIndexes[j]] = nj[j];
	for(uint e = 0; e < Ne; ++e) n[eIndexes[e]] = problem.ne[e];
	
	return diagnostic;
}

void EquilibriumSolver::UnderRelax(VectorXd& n, const VectorXd& nOld) const
{
	for(uint i = 0; i < uint(n.rows()); ++i)
	{
		const double delta  = n[i] - nOld[i];
		const double factor = max(1.0, -2.0 * delta/nOld[i]);
		
		n[i] = nOld[i] + delta/factor;
	}
}
